package com.husd.leetcode.tree;

import com.husd.leetcode.entry.ColorEnum;
import com.husd.leetcode.entry.TreeNode;

/**
 * 红黑树
 *
 * @author hushengdong
 */
public class RedBlackTree implements Tree {

    private final static int INSERT_CASE1 = 1;
    private final static int INSERT_CASE2 = 2;
    private final static int INSERT_CASE3 = 3;
    private final static int INSERT_CASE_END = 0;

    //final TreeNode nil = new TreeNode(null, ColorEnum.BLACK);
    TreeNode root;
    int size = 0;

    public RedBlackTree(TreeNode root) {
        this.root = root;
    }

    public RedBlackTree() {
        this.root = null;
    }

    @Override
    public void insert(int z) {

        System.out.println("insert " + z);
        //y means pre
        TreeNode y = null;
        TreeNode x = root;

        while (x != null) {

            y = x;
            if (z < x.value) {
                x = x.left;
            } else if (z > x.value) {
                x = x.right;
            } else if (z == x.value) {
                return;
            }
        }

        //默认是红色节点
        TreeNode newNode = new TreeNode(z, ColorEnum.RED);
        //红黑树，我新加了一个pre属性，方便操作。
        newNode.pre = y;

        if (y == null) {
            //表示空的树
            root = newNode;
        } else if (z < y.value) {
            y.left = newNode;
        } else if (z > y.value) {
            y.right = newNode;
        }

        //到这里为止，和普通的二叉搜索树，还没有什么区别。
        //接下来就需要判断下颜色，然后看看是不是不满足条件了，然后修复。
        insertColorFix2(newNode);
    }

    private void insertColorFix2(TreeNode z) {

        //根节点单独处理
        if (rootNode(z)) {
            z.setBlack();
            return;
        }

        TreeNode focus = z;

        TreeNode father = focus.pre;
        //父节点是黑色的，什么都不做。
        if (father.isBlack()) {
            return;
        }
        //
        TreeNode uncle = uncleNode(focus);
        int c = calcCase(focus, father, uncle);
        while (c != INSERT_CASE_END) {

            if (c == INSERT_CASE1) {
                focus = insertCase1(father, uncle);

                father = focus.pre;
                uncle = uncleNode(focus);
                c = calcCase(focus, father, uncle);
            } else if (c == INSERT_CASE2) {

                focus = insertCase2(father, uncle);
                father = focus.pre;
                uncle = uncleNode(focus);
                c = calcCase(focus, father, uncle);
            } else if (c == INSERT_CASE3) {

                insertCase3(father, brotherNode(focus));
                c = INSERT_CASE_END;
            }
        }
        //根节点要设置黑色
        root.setBlack();
    }

    private int calcCase(TreeNode focus, TreeNode father, TreeNode uncle) {

        if (TreeUtils.red(uncle)) {
            return INSERT_CASE1;
        }
        if (father == null || father.pre == null) {
            //叔叔节点是null，默认为黑色。
            return INSERT_CASE_END;
        }
        boolean leftN = leftNode(focus, father);
        if (leftN) {
            return INSERT_CASE3;
        }
        return INSERT_CASE2;
    }


    private void insertCase3(TreeNode father, TreeNode brotherNode) {

        TreeNode fatherFather = father.pre;
        rightRotate(fatherFather);
        //父节点和兄弟节点互换颜色
        ColorEnum fatherColor = getColorEnum(father);
        ColorEnum brotherColor = getColorEnum(brotherNode);

        if (father != null) {
            father.color = brotherColor;
        }
        if (brotherNode != null) {
            brotherNode.color = fatherColor;
        }
    }

    private ColorEnum getColorEnum(TreeNode node) {

        return node == null ? ColorEnum.BLACK : node.color;
    }

    private TreeNode insertCase2(TreeNode father, TreeNode uncle) {

        leftRotate(father);
        return father;
    }

    private TreeNode insertCase1(TreeNode father, TreeNode uncle) {

        //将关注节点 a 的父节点 b、叔叔节点 d 的颜色都设置成黑色
        father.setBlack();
        uncle.setBlack();
        //将关注节点 a 的祖父节点 c 的颜色设置成红色
        TreeNode fatherFather = father.pre;
        fatherFather.setRed();
        //关注节点变成 a 的祖父节点 c
        return fatherFather;
    }

    private boolean rootNode(TreeNode z) {

        return z == root;
    }

    private TreeNode uncleNode(TreeNode node) {

        TreeNode father = node.pre;
        if (father == null) return null;
        TreeNode fatherfather = father.pre;
        if (fatherfather == null) return null;
        boolean left = father.value < fatherfather.value;
        return left ? fatherfather.right : fatherfather.left;
    }

    private TreeNode brotherNode(TreeNode node) {

        TreeNode father = node.pre;
        if (father == null) return null;
        boolean left = node.value < father.value;
        return left ? father.right : father.left;
    }

    private boolean leftNode(TreeNode son, TreeNode father) {

        return son.value.compareTo(father.value) < 0;
    }

    private boolean rightNode(TreeNode son, TreeNode father) {

        return son.value.compareTo(father.value) > 0;
    }


    //这个算法导论的版本。目前没测试通过
    public void insert2(int z) {

        //y means pre
        TreeNode y = null;
        TreeNode x = root;

        while (x != null) {

            y = x;
            if (z < x.value) {
                x = x.left;
            } else if (z > x.value) {
                x = x.right;
            } else if (z == x.value) {
                return;
            }
        }

        //默认是红色节点
        TreeNode newNode = new TreeNode(z, ColorEnum.RED);
        //红黑树，我新加了一个pre属性，方便操作。
        newNode.pre = y;

        if (y == null) {
            //表示空的树
            root = newNode;
        } else if (z < y.value) {
            y.left = newNode;
        } else if (z > y.value) {
            y.right = newNode;
        }

        //到这里为止，和普通的二叉搜索树，还没有什么区别。
        //接下来就需要判断下颜色，然后看看是不是不满足条件了，然后修复。
        insertColorFix(newNode);
    }

    private void insertColorFix(TreeNode z) {

        if (root == null) return;
        if (root.left == null && root.right == null) {
            root.setBlack();
            return;
        }
        TreeNode y = null;
        //如果新节点的父节点是黑色的，插入一个红色的，不会破坏平衡。
        while (z.pre != null && z.pre.isRed()) {
            //z.pre.pre may be null
            if (z.pre.pre != null && z.pre == z.pre.pre.left) {
                y = z.pre.pre.right;
                if (y.isRed()) {
                    z.pre.setBlack();
                    y.setBlack();
                    z.pre.pre.setRed();
                    z = z.pre.pre;
                } else {
                    if (z == z.pre.right) {
                        z = z.pre;
                        leftRotate(z);
                        z.pre.setBlack();
                        z.pre.pre.setRed();
                        rightRotate(z);
                    } else {
                        z = z.pre;
                        rightRotate(z);
                        z.pre.setBlack();
                        z.pre.pre.setRed();
                        leftRotate(z);
                    }
                }
            }
        }
        root.setBlack();
    }

    public void rightRotate(TreeNode x) {

        if (x == null) return;
        TreeNode y = x.left;
        if (y == null) return;
        x.left = y.right;   // turn y's right subtree into x's left subtree
        if (y.right != null) {
            y.right.pre = x;
        }
        y.pre = x.pre;
        if (x.pre == null) {
            root = y;
        } else if (x == x.pre.right) {
            x.pre.right = y;
        } else {
            x.pre.left = y;
        }
        y.right = x; // put x on y's right
        x.pre = y;
    }

    public void leftRotate(TreeNode x) {

        if (x == null) return;
        TreeNode y = x.right;
        if (y == null) return;
        x.right = y.left;   // turn y's left subtree into x's right subtree
        if (y.left != null) {
            y.left.pre = x;
        }
        y.pre = x.pre;
        if (x.pre == null) {
            root = y;
        } else if (x == x.pre.left) {
            x.pre.left = y;
        } else {
            x.pre.right = y;
        }
        y.left = x; // put x on y's left
        x.pre = y;
    }

    @Override
    public void delete(int value) {

    }

    //移植代码
    private void transplant(TreeNode u, TreeNode v) {

        if (u.pre == null) {
            root = v;
        } else if (u == u.pre.left) {
            u.pre.left = v;
        } else if (u.pre.right == v) {
            v.pre = u.pre;
        }
    }

    @Override
    public TreeNode find(TreeNode root, int value) {

        return null;
    }

    @Override
    public int depth(TreeNode root) {
        return 0;
    }
}




